import numpy as np
import gym
import time
import pygame
from marioEnv import create_train_env, ACTIONS
from nes_py_v1.nes_env import GameStateManager

env = create_train_env((1, 1), process_frame=False, skip=False, skip_cnt=4)
state = env.reset()

print("ram size: ", env.ram.shape)

pygame.init()
pygame.display.set_mode((256, 240))

captures = []
rams = []
actions = []

#data = np.load('database.npz')
#old_capture = data['capture']
#old_rams = data['ram']
#old_actions = data['actions']
#captures.append(old_capture)
#rams.append(old_rams)
#actions.append(old_actions)

tmp_captures = []
tmp_rams = []
tmp_actions = []

state_mgr = GameStateManager(5)
saved = False
saved_state = None

while True:
	env.render()
	# get input
	key_pressed = pygame.key.get_pressed()
	action = 0
	if key_pressed[pygame.K_a]:
		if key_pressed[pygame.K_SPACE]: 
			action = 2 
		else: 
			action = 1
	elif key_pressed[pygame.K_d]:
		if key_pressed[pygame.K_SPACE]: 
			action = 4 
		else: 
			action = 3
	elif key_pressed[pygame.K_SPACE]:
		action = 5
	if key_pressed[pygame.K_j] and not saved:
		saved_state = state_mgr.save(env)
		saved = True
		print('saved')
	if key_pressed[pygame.K_k] and saved:
		state_mgr.load(env, saved_state)
		state_mgr.release(saved_state)
		saved = False
		saved_state = None
		print('load')
	if key_pressed[pygame.K_ESCAPE]:
		pygame.quit()
		break
	
	# collect data
	tmp_captures.append(state[np.newaxis,:])
	tmp_rams.append(env.ram[np.newaxis,:])
	tmp_actions.append(action)
	# step
	state, reward, done, info = env.step(action)
	if done:
		print('done')
		state = env.reset()
		if info['flag_get']:
			captures.extend(tmp_captures)
			rams.extend(tmp_rams)
			actions.extend(tmp_actions)
		tmp_captures = []
		tmp_rams = []
		tmp_actions = []
	time.sleep(1./100.)
	
# database_captures = np.concatenate(captures, axis=0)
# print(database_captures.shape)
# database_rams = np.concatenate(rams, axis=0)
# print(database_rams.shape)
# database_actions = np.array(actions)
# print(database_actions.shape)

# np.savez('database.npz', captures=database_captures, rams=database_rams, actions=database_actions)
# print('complete!')